#
# This file is a part of the open-eBackup project.
# This Source Code Form is subject to the terms of the Mozilla Public License, v. 2.0.
# If a copy of the MPL was not distributed with this file, You can obtain one at
# http://mozilla.org/MPL/2.0/.
#
# Copyright (c) [2024] Huawei Technologies Co.,Ltd.
#
# THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND,
# EITHER EXPRESS OR IMPLIED, INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT,
# MERCHANTABILITY OR FIT FOR A PARTICULAR PURPOSE.
#

import json
import os
import sys

import pexpect

from antdb.common.antdb_common import get_version
from antdb.common.antdb_exec_sql import ExecAntDBSql
from antdb.common.const import PexpectResult
from antdb.common.error_code import ErrorCode
from antdb.common.util.antdb_common_utils import AntDBCommonUtils
from common import cleaner
from common.cleaner import clear
from common.common import exter_attack, execute_cmd, check_command_injection
from common.common_models import ActionResult, QueryClusterResponse, ClusterNodeInfo, ClusterNodeExtendInfo
from common.const import ParamConstant, ExecuteResultEnum, AuthType, SysData, CMDResult, RoleType
from common.logger import Logger
from common.parse_parafile import get_env_variable, ParamFileUtil
from common.util import check_utils, check_user_utils
from common.util.check_utils import is_valid_id
from common.util.cmd_utils import cmd_format
from common.util.exec_utils import exec_overwrite_file, check_path_valid

LOGGER = Logger().get_logger("antdb.log")


class ClusterNodesChecker:
    def __init__(self, check_type, request_pid):
        self.check_type = check_type
        self.pid = request_pid
        self.context = ParamFileUtil.parse_param_file(self.pid)
        self.application = self.context.get("application", {})
        self.env = self.context.get("appEnv", {})
        self.extend_info = self.application.get("extendInfo", {})
        self.env_extend = self.env.get("extendInfo", {})
        self.os_username = self.extend_info.get("osUsername")
        self.adbhamgr_path = self.extend_info.get("adbhamgrPath", "/etc/adbhamgr/adbhamgr_antdbcluster.yaml")
        self.rep_user = self.application.get("auth", {}).get("extendInfo", {}).get("dbStreamRepUser", "repl")
        self.service_ip = self.extend_info.get("serviceIp")
        self.inst_port = self.extend_info.get("instancePort")
        self.result_file = os.path.join(ParamConstant.RESULT_PATH, f"result{self.pid}")
        self.auth_type = get_env_variable(f"application_auth_authType_{self.pid}")
        self.nodes = self.env_extend.get("allNodes")
        self.port = self.extend_info.get("instancePort")
        self.client_path = os.path.realpath(os.path.join(self.extend_info.get("clientPath"), "bin", "adb"))
        self.pg_version = None
        self.antdb_version = None

    @staticmethod
    def get_nodes(os_username, adbhamgr_config):
        if check_command_injection(adbhamgr_config):
            LOGGER.error(f"adbhamgr config:{adbhamgr_config} check error!")
            return []
        cmd = cmd_format("su - {} -c 'adbhamgrctl -c {} list'", os_username, adbhamgr_config)
        ret_code, std_out, std_err = execute_cmd(cmd)
        if ret_code != CMDResult.SUCCESS.value or not std_out:
            LOGGER.error(f"get_nodes out: {std_out},error: {std_err}")
            return []
        nodes = []
        rows = std_out.splitlines()
        rows_len = len(rows)

        for i in range(3, rows_len - 1):
            # | adbhamgr-2 | 8.40.160.161:6655 | Leader       | running |  1 |           |
            row = rows[i]
            columns = row.split("|")
            if len(columns) <= 6:
                LOGGER.error("Failed to get nodes: number of columns is incorrect")
                return []
            role = columns[3].strip()
            status = columns[4].strip()
            host = columns[2].strip().split(":")
            hostname = host[0]
            node = {
                'hostname': hostname,
                'role': str(RoleType.PRIMARY.value) if role in ['Leader'] else str(RoleType.STANDBY.value),
                'status': status
            }
            nodes.append(node)
        return nodes

    @exter_attack
    def check_cluster_nodes(self):
        LOGGER.info(f"Begin to check cluster by os param: {self.context}, pid: {self.pid}")
        param = ActionResult(code=ExecuteResultEnum.INTERNAL_ERROR, bodyErr=ErrorCode.CHECK_CLUSTER_FAILED,
                             message="check cluster failed!")
        try:
            param = self._check_cluster_nodes()
        except Exception as e:
            param = ActionResult(code=ExecuteResultEnum.INTERNAL_ERROR, bodyErr=ErrorCode.CHECK_CLUSTER_FAILED,
                                 message="An exception occur when check cluster failed.")
            LOGGER.error(f"Failed to check cluster!pid: {self.pid}", e)
        finally:
            cleaner.clear(SysData.SYS_STDIN)
            LOGGER.info('Clearing data successfully')
            LOGGER.info(f"Finally to check cluster, pid: {self.pid}, param:{param}")
            exec_overwrite_file(self.result_file, param.dict(by_alias=True))

    def check_application(self, nodes):
        LOGGER.info(f"Already get nodes: {nodes}, pid: {self.pid}")
        result = [
            f"{AntDBCommonUtils.domain_2_ip(i.get('hostname'))}:{i.get('port')}"
            for i in nodes
            if i.get('hostname') and i.get('port')
        ]
        LOGGER.info(f"Already get nodes: {result}, pid: {self.pid}")
        node_list = self.nodes.split(",")
        if len(result) != len(node_list):
            return ActionResult(code=ExecuteResultEnum.INTERNAL_ERROR, bodyErr=ErrorCode.CHECK_CLUSTER_FAILED,
                                message="Missing or redundant cluster nodes.")
        res = set(result)
        node_list = set(node_list)
        LOGGER.info(f"Already get result: {res}, node_list: {node_list}, pid: {self.pid}")
        if not res.difference(node_list):
            LOGGER.info(f"Success to check cluster pid: {self.pid}")
            return ActionResult(code=ExecuteResultEnum.SUCCESS, bodyErr=ExecuteResultEnum.SUCCESS,
                                message=json.dumps({"pg_version": self.pg_version,
                                                    "antdb_version": self.antdb_version}))
        param = ActionResult(code=ExecuteResultEnum.INTERNAL_ERROR, bodyErr=ErrorCode.CHECK_CLUSTER_FAILED,
                             message="Nodes not belong to a cluster.")
        return param

    def query_app_cluster(self, nodes):
        result = QueryClusterResponse(name=self.env.get('name'), subType='AntDBClusterInstance')
        for i in nodes:
            tmp_extend = ClusterNodeExtendInfo()
            tmp_extend.role = 1 if i.get('role') in ['primary', 'master', 'Leader'] else 2
            tmp_extend.status = 1 if i.get('status') == 'running' else 0
            tmp_node = ClusterNodeInfo(endpoint=AntDBCommonUtils.domain_2_ip(i.get('hostname')),
                                       subType=self.application.get('subType'), extendInfo=tmp_extend)
            result.nodes.append(tmp_node)
        LOGGER.info(f"Success to check cluster all_info: {result.dict(by_alias=True)}")
        return result

    def check_cluster_by_request_type(self, nodes):
        if self.check_type == "CheckApplication":
            return self.check_application(nodes)
        if self.check_type == "QueryCluster":
            return self.query_app_cluster(nodes)
        LOGGER.error(f"Unknown command: {self.check_type} from command line!")
        return ActionResult(code=ExecuteResultEnum.INTERNAL_ERROR, bodyErr=ErrorCode.GET_CLUSTER_NODE_FAILED,
                            message=f"check_type is not exist!")

    def login_database_and_execute_cmd(self, cmd, pwd=None):
        result, child = self._checkout_os_user()
        if not result:
            return False, child

        child.sendline(cmd)
        index = child.expect(PexpectResult.DB_LOGIN_PASSWORD)
        if index in (0, 1):
            LOGGER.error(f"Login denied: {self.os_username}, pid: {self.pid}")
            result = ActionResult(code=ExecuteResultEnum.INTERNAL_ERROR, bodyErr=ErrorCode.LOGIN_FAILED,
                                  message=f"Login denied")
            child.close()
            return False, result

        if not pwd:
            pwd = get_env_variable(f"application_auth_authPwd_{self.pid}")
        child.sendline(pwd)
        clear(pwd)
        LOGGER.info(f"Input password, pid: {self.pid}")
        index = child.expect(PexpectResult.HTML_RESULT)
        LOGGER.info(f"Get Node Info,pool:{child.before} index:{index},pid: {self.pid}")
        if index in (0, 1):
            LOGGER.error(f"Get pg pool info failed: {self.os_username}, pid: {self.pid}")
            result = ActionResult(code=ExecuteResultEnum.INTERNAL_ERROR, bodyErr=ErrorCode.GET_CLUSTER_NODE_FAILED,
                                  message=f"Get pg pool info failed!")
            child.close()
            return False, result
        return True, child

    def check_repl_user(self):
        LOGGER.info("Begin to check repl user!")
        cmd = f"{self.client_path} \"dbname=postgres replication=true\" -c \"IDENTIFY_SYSTEM;\" " \
              f"-U {self.rep_user} -h {self.service_ip} -H -W"
        repl_pwd = None
        child = None
        try:
            repl_pwd = get_env_variable(f"application_auth_extendInfo_dbStreamRepPwd_{self.pid}")
            result, child = self.login_database_and_execute_cmd(cmd, repl_pwd)
            if not result:
                LOGGER.error("Check repl user failed!")
                return False
            LOGGER.info("Check repl user success!")
            return True
        finally:
            if child and not isinstance(child, ActionResult):
                child.close()
            if repl_pwd:
                clear(repl_pwd)

    def get_nodes_domain_and_port_by_database_patroni(self):
        LOGGER.info("Begin to check patroni client!")
        cmd = cmd_format("su - {} -c 'adbhamgrctl -c {} list'", self.os_username, self.adbhamgr_path)
        ret_code, std_out, std_err = execute_cmd(cmd)
        param = ActionResult(code=ExecuteResultEnum.INTERNAL_ERROR, bodyErr=ErrorCode.GET_CLUSTER_NODE_FAILED,
                             message=f"check_type is not exist!")
        if ret_code != CMDResult.SUCCESS.value:
            LOGGER.info(f"get_nodes_domain_and_port_by_database_patroni, out: {std_out}, error: {std_err}")
            return param
        if len(std_out) <= 0:
            return param
        rows = std_out.splitlines()
        nodes = []
        rows_len = len(rows)
        node_list = self.nodes.split(",")
        node_list = set(node_list)
        for i in range(3, rows_len - 1):
            row = rows[i]
            columns = row.split("|")
            role = columns[3].strip()
            status = columns[4].strip()
            host = columns[2].strip().split(":")
            hostname = host[0]
            port = host[1]
            node = {'hostname': hostname, 'role': role, 'port': port, 'status': status}
            nodes.append(node)
        return self.check_cluster_by_request_type(nodes)

    def _check_cluster_nodes_first(self):
        for check in (self.os_username, self.client_path):
            if not AntDBCommonUtils.check_special_characters(check):
                return ActionResult(code=ExecuteResultEnum.INTERNAL_ERROR, bodyErr=ErrorCode.CHECK_CONNECTIVITY_FAILED,
                                    message=f"String contains special characters!")
        if not AntDBCommonUtils.check_db_user_valid(self.rep_user):
            return ActionResult(code=ExecuteResultEnum.INTERNAL_ERROR, bodyErr=ErrorCode.CHECK_CONNECTIVITY_FAILED,
                                message=f"Repl user name is invalid!")
        if not AntDBCommonUtils.check_path(self.client_path):
            LOGGER.info(f"{self.client_path} is invalid.")
            return ActionResult(code=ExecuteResultEnum.INTERNAL_ERROR,
                                bodyErr=ErrorCode.DATABASE_INSTALLATION_PATH_IS_INVALID,
                                message=f"Database installation path is invalid!")
        if not AntDBCommonUtils.check_path(self.adbhamgr_path) or not os.path.isfile(self.adbhamgr_path):
            LOGGER.info(f"{self.adbhamgr_path} is not a file.")
            return ActionResult(code=ExecuteResultEnum.INTERNAL_ERROR,
                                bodyErr=ErrorCode.FULL_PATH_OF_PATRONI_CONFIGURATION_FILE_IS_INVALID,
                                message=f"The full path of Patroni configuration file is invalid!")
        if not check_utils.is_ip_address(self.service_ip) or self.service_ip not in AntDBCommonUtils.get_local_ips():
            LOGGER.info(f"{self.service_ip} is invalid.")
            return ActionResult(code=ExecuteResultEnum.INTERNAL_ERROR, bodyErr=ErrorCode.SERVICE_IP_IS_INVALID,
                                message=f"Service ip is invalid or can't be localhost!")
        if not AntDBCommonUtils.check_os_user(self.os_username, self.client_path)[0]:
            return ActionResult(code=ExecuteResultEnum.INTERNAL_ERROR, bodyErr=ErrorCode.USER_IS_NOT_EXIST,
                                message=f"Os username is not exist!")
        return ActionResult(code=ExecuteResultEnum.SUCCESS, bodyErr=ExecuteResultEnum.SUCCESS,
                            message='check success')

    def _check_cluster_nodes(self):
        check_first_result = self._check_cluster_nodes_first()
        if check_first_result.code != ExecuteResultEnum.SUCCESS:
            return check_first_result
        client_path = self.extend_info.get("clientPath")
        get_version_res, pg_version, antdb_version = get_version(
            self.pid, client_path, self.os_username, self.service_ip, self.port)
        if not get_version_res:
            return ActionResult(code=ExecuteResultEnum.INTERNAL_ERROR, bodyErr=ErrorCode.GET_VERSION_FAILED,
                                message="Get version failed!")
        self.pg_version = pg_version
        self.antdb_version = antdb_version
        if int(self.auth_type) == AuthType.APP_PASSWORD.value:
            # 检查用户名和密码,通过用户名和密码查询归档开关是否开启
            return self.get_nodes_domain_and_port_by_database_patroni()
        return ActionResult(code=ExecuteResultEnum.INTERNAL_ERROR, bodyErr=ErrorCode.LOGIN_FAILED,
                            message=f"Login denied")

    def _query_archive_mode(self, db_user_name, db_pwd):
        sql_cmd = "show archive_mode;"
        pg_sql = ExecAntDBSql(self.pid, self.extend_info.get("clientPath"), self.service_ip, self.port)
        return_code, _, _ = pg_sql.exec_sql_cmd_resource(self.os_username, sql_cmd, db_user_name, db_pwd)
        return return_code

    def _checkout_os_user(self):
        # 登录
        child = pexpect.spawn(f"su - {self.os_username}", encoding="utf-8", timeout=20)
        LOGGER.info(f"Success to execute su - {self.os_username}, pid: {self.pid}")
        index = child.expect(PexpectResult.OS_LOGIN_RESULT)

        if index in (0, 1):
            LOGGER.error(
                f"Checkout user: {self.os_username} failed!pid: {self.pid}")
            child.close()
            param = ActionResult(code=ExecuteResultEnum.INTERNAL_ERROR, bodyErr=ErrorCode.LOGIN_FAILED,
                                 message="Login failed!")
            return False, param
        LOGGER.info(f"Success check os user: {self.os_username}, pid: {self.pid}")
        return True, child


if __name__ == '__main__':
    function = sys.argv[1]
    pid = sys.argv[2]  # argv
    for line in sys.stdin:
        SysData.SYS_STDIN = line
        break
    if not is_valid_id(pid):
        LOGGER.warn(f'pid is invalid!')
        sys.exit(1)
    LOGGER.info("Begin to check Cluster!")
    checker = ClusterNodesChecker(check_type=function, request_pid=pid)
    checker.check_cluster_nodes()
